#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
优化版RAG系统 - 基于优化的文档解析器
=====================================

这是原始 rag1.py 的优化版本，集成了先进的文档解析功能。
主要改进：
1. 使用优化的文档解析器
2. 更好的错误处理
3. 更丰富的内容提取
4. 更高的处理性能

作者：AI Assistant
基于：原始 rag1.py
版本：2.0.0 (优化版)
"""

import chromadb
from chromadb.utils import embedding_functions
import os
import logging
from openai import OpenAI
from typing import List, Dict, Any, Optional

# 导入优化的文档解析器
try:
    from optimized_document_parser import (
        DocumentConfig,
        load_employee_file_enhanced,
        load_employee_file_simple,
        create_rag_friendly_chunks
    )
    OPTIMIZED_PARSER_AVAILABLE = True
    print("✅ 成功导入优化文档解析器")
except ImportError as e:
    OPTIMIZED_PARSER_AVAILABLE = False
    print(f"⚠️  优化解析器不可用，回退到基础实现: {str(e)}")
    from docx import Document

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 初始化ChromaDB客户端
chroma_client = chromadb.PersistentClient(path="./chroma_db")


def load_employee_file(file_path: str) -> str:
    """
    优化版员工手册加载函数
    
    这是对原始函数的直接替换，保持API兼容性的同时提供更强大的功能。
    
    Args:
        file_path: 员工手册文件路径
        
    Returns:
        解析后的文档文本内容
        
    Raises:
        Exception: 当文档解析失败时
    """
    if OPTIMIZED_PARSER_AVAILABLE:
        try:
            # 使用优化解析器
            logger.info(f"使用优化解析器处理文档: {file_path}")
            
            # 配置解析选项
            config = DocumentConfig(
                extract_paragraphs=True,
                extract_tables=True,          # 新增：提取表格内容
                extract_images=False,         # 对于RAG，图片信息通常不需要
                clean_text=True,             # 新增：清理文本
                remove_empty_paragraphs=True, # 新增：移除空段落
                preserve_formatting=False,    # RAG通常不需要格式信息
                verbose=False
            )
            
            result = load_employee_file_enhanced(file_path, config)
            
            if result.success:
                text_content = result.to_text()
                
                # 记录处理统计信息
                stats = result.get_statistics()
                logger.info(f"文档解析完成: {len(result.elements)}个元素, "
                          f"处理时间: {result.processing_time:.2f}秒")
                logger.info(f"元素统计: {stats}")
                
                # 如果有警告，记录但不影响处理
                if result.warnings:
                    for warning in result.warnings:
                        logger.warning(f"解析警告: {warning}")
                
                return text_content
            else:
                # 解析失败，记录错误并抛出异常
                error_msg = f"优化解析器处理失败: {'; '.join(result.errors)}"
                logger.error(error_msg)
                raise Exception(error_msg)
                
        except Exception as e:
            logger.error(f"优化解析器异常: {str(e)}")
            # 可以选择回退到原始实现或直接抛出异常
            raise Exception(f"文档解析失败: {str(e)}")
    else:
        # 回退到原始实现
        logger.warning("使用原始解析器（功能受限）")
        return _original_load_employee_file(file_path)


def load_employee_file_advanced(file_path: str, 
                               include_tables: bool = True,
                               include_metadata: bool = False) -> Dict[str, Any]:
    """
    高级版员工手册加载函数
    
    提供更多选项和详细信息，适用于需要更多控制的场景。
    
    Args:
        file_path: 员工手册文件路径
        include_tables: 是否包含表格内容
        include_metadata: 是否返回元数据信息
        
    Returns:
        包含文本内容和可选元数据的字典
    """
    if not OPTIMIZED_PARSER_AVAILABLE:
        # 回退实现
        text = _original_load_employee_file(file_path)
        return {
            'text': text,
            'metadata': {'parser': 'original', 'length': len(text)}
        }
    
    config = DocumentConfig(
        extract_paragraphs=True,
        extract_tables=include_tables,
        extract_images=False,
        clean_text=True,
        remove_empty_paragraphs=True,
        verbose=True
    )
    
    result = load_employee_file_enhanced(file_path, config)
    
    if result.success:
        response = {
            'text': result.to_text(),
            'success': True,
            'processing_time': result.processing_time
        }
        
        if include_metadata:
            response.update({
                'metadata': result.metadata,
                'statistics': result.get_statistics(),
                'elements': len(result.elements),
                'warnings': result.warnings
            })
        
        return response
    else:
        raise Exception(f"高级解析失败: {'; '.join(result.errors)}")


def _original_load_employee_file(file_path: str) -> str:
    """
    原始的员工手册加载函数（备用）
    
    保留原始实现作为备用方案
    """
    try:
        document = Document(file_path)
        all_text = []
        for para in document.paragraphs:
            clean_text = para.text.strip()
            all_text.append(clean_text)
        return '\n'.join(all_text)
    except Exception as e:
        raise Exception(f"原始解析器失败: {str(e)}")


def split_text_intelligent(text: str, chunk_size: int = 200) -> List[str]:
    """
    智能文本切分函数 - 替换原有的简单切分
    
    使用优化解析器的智能切分功能，如果不可用则回退到原始实现。
    
    Args:
        text: 待切分的文本
        chunk_size: 块大小
        
    Returns:
        切分后的文本块列表
    """
    if OPTIMIZED_PARSER_AVAILABLE:
        try:
            # 如果有完整的文件路径，使用create_rag_friendly_chunks
            # 这里我们假设text是从文件解析来的，使用简单的切分逻辑
            return _split_text_with_structure(text, chunk_size)
        except Exception as e:
            logger.warning(f"智能切分失败，使用原始切分: {str(e)}")
            return _original_split_text(text, chunk_size)
    else:
        return _original_split_text(text, chunk_size)


def _split_text_with_structure(text: str, chunk_size: int) -> List[str]:
    """
    带结构的文本切分
    
    尝试保持文档结构的完整性
    """
    import re
    
    # 按段落分割
    paragraphs = text.split('\n\n')
    chunks = []
    current_chunk = ""
    
    for paragraph in paragraphs:
        paragraph = paragraph.strip()
        if not paragraph:
            continue
        
        # 如果当前块加上新段落超过大小限制
        if len(current_chunk) + len(paragraph) > chunk_size and current_chunk:
            chunks.append(current_chunk.strip())
            current_chunk = paragraph
        else:
            if current_chunk:
                current_chunk += "\n\n" + paragraph
            else:
                current_chunk = paragraph
    
    # 添加最后一个块
    if current_chunk.strip():
        chunks.append(current_chunk.strip())
    
    # 对过长的块进行二次切分
    final_chunks = []
    for chunk in chunks:
        if len(chunk) <= chunk_size:
            final_chunks.append(chunk)
        else:
            # 按句子切分
            sentences = re.split(r'[。！？\n]', chunk)
            sub_chunk = ""
            
            for sentence in sentences:
                sentence = sentence.strip()
                if not sentence:
                    continue
                
                if len(sub_chunk) + len(sentence) > chunk_size and sub_chunk:
                    final_chunks.append(sub_chunk.strip())
                    sub_chunk = sentence
                else:
                    if sub_chunk:
                        sub_chunk += "。" + sentence
                    else:
                        sub_chunk = sentence
            
            if sub_chunk.strip():
                final_chunks.append(sub_chunk.strip())
    
    return final_chunks


def _original_split_text(text: str, chunk_size: int) -> List[str]:
    """原始的文本切分函数"""
    chunks = []
    for i in range(0, len(text), chunk_size):
        chunks.append(text[i:i + chunk_size])
    return chunks


def create_rag_chunks_from_file(file_path: str, chunk_size: int = 500) -> List[str]:
    """
    从文件直接创建RAG友好的文档块
    
    这是一个新增的便利函数，专门为RAG系统优化
    
    Args:
        file_path: 员工手册文件路径
        chunk_size: 块大小
        
    Returns:
        RAG友好的文档块列表
    """
    if OPTIMIZED_PARSER_AVAILABLE:
        try:
            chunks = create_rag_friendly_chunks(file_path, chunk_size=chunk_size)
            logger.info(f"使用优化RAG切分: 生成{len(chunks)}个块")
            return chunks
        except Exception as e:
            logger.warning(f"优化RAG切分失败: {str(e)}")
    
    # 回退到传统方法
    logger.info("使用传统RAG切分方法")
    text = load_employee_file(file_path)
    return split_text_intelligent(text, chunk_size)


def text_embeddings(texts: List[str]) -> List[List[float]]:
    """
    文本向量化函数 - 保持原有接口
    
    这个函数保持不变，但可以处理优化解析器生成的更高质量文本
    """
    client = OpenAI(
        api_key=os.getenv("DASHSCOPE_API_KEY"),
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    if isinstance(texts, str):
        texts = [texts]

    batch_size = 10
    embeddings = []

    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        completion = client.embeddings.create(
            model="text-embedding-v4",
            input=batch,
            dimensions=1024,
            encoding_format="float"
        )

        batch_embeddings = [item.embedding for item in completion.data]
        embeddings.extend(batch_embeddings)

    return embeddings


def store_in_chroma_optimized(texts: List[str], collection_name: str = "employee_manual") -> bool:
    """
    优化的ChromaDB存储函数
    
    增加了错误处理和进度跟踪
    
    Args:
        texts: 文本列表
        collection_name: 集合名称
        
    Returns:
        是否存储成功
    """
    try:
        logger.info(f"开始存储{len(texts)}个文档块到ChromaDB")
        
        # 计算嵌入向量
        embeddings = text_embeddings(texts)
        
        # 创建嵌入函数
        embedding_function = embedding_functions.OpenAIEmbeddingFunction(
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            model_name="text-embedding-v4",
            api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )

        # 尝试删除已存在的集合
        try:
            chroma_client.delete_collection(name=collection_name)
            logger.info(f"删除已存在的集合: {collection_name}")
        except Exception:
            pass  # 集合不存在，忽略错误

        # 创建新集合
        collection = chroma_client.create_collection(
            name=collection_name,
            embedding_function=embedding_function
        )

        # 批量添加文档
        batch_size = 50  # 批处理大小
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]
            batch_embeddings = embeddings[i:i + batch_size]
            batch_ids = [f"text_{j}" for j in range(i, i + len(batch_texts))]
            
            collection.add(
                documents=batch_texts,
                embeddings=batch_embeddings,
                ids=batch_ids
            )
            
            logger.info(f"已存储 {min(i + batch_size, len(texts))}/{len(texts)} 个文档块")

        logger.info("ChromaDB存储完成")
        return True
        
    except Exception as e:
        logger.error(f"ChromaDB存储失败: {str(e)}")
        return False


def rag_answer_question_enhanced(question: str, 
                                top_k: int = 3,
                                collection_name: str = "employee_manual") -> Dict[str, Any]:
    """
    增强版RAG问答函数
    
    提供更详细的结果和更好的错误处理
    
    Args:
        question: 用户问题
        top_k: 检索文档数量
        collection_name: 集合名称
        
    Returns:
        包含答案和详细信息的字典
    """
    try:
        logger.info(f"处理问题: {question}")
        
        # 创建嵌入函数
        embedding_function = embedding_functions.OpenAIEmbeddingFunction(
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            model_name="text-embedding-v4",
            api_base="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )

        # 获取集合
        collection = chroma_client.get_collection(
            name=collection_name,
            embedding_function=embedding_function
        )

        # 检索相关文档
        results = collection.query(
            query_texts=[question],
            n_results=top_k
        )

        if not results["documents"] or not results["documents"][0]:
            return {
                'success': False,
                'error': '未找到相关文档',
                'answer': '抱歉，没有找到相关信息来回答您的问题。'
            }

        retrieved_chunks = list(set(results["documents"][0]))
        distances = results["distances"][0] if results["distances"] else []
        
        logger.info(f"检索到 {len(retrieved_chunks)} 个相关文档块")

        # 构建上下文
        context = "\n".join([f"- {chunk}" for chunk in retrieved_chunks])
        
        # 构建优化的提示词
        prompt = f"""你是一个专业的HR助理，请仔细阅读以下员工手册内容并准确回答用户问题。

【员工手册内容】：
{context}

【用户问题】：
{question}

【回答要求】：
1. 严格基于员工手册内容回答，不要编造或推测
2. 如果手册中没有相关信息，请明确说明"手册中未提及此信息"
3. 回答应简洁明了，条理清晰
4. 如有多个要点，请分点列出
5. 如果涉及具体条款，可引用相关内容

请回答："""

        # 调用LLM生成回答
        response = get_response([{"role": "user", "content": prompt}])
        
        return {
            'success': True,
            'answer': response,
            'question': question,
            'retrieved_chunks': retrieved_chunks,
            'chunk_count': len(retrieved_chunks),
            'distances': distances,
            'context_length': len(context)
        }
        
    except Exception as e:
        error_msg = f"RAG问答处理失败: {str(e)}"
        logger.error(error_msg)
        return {
            'success': False,
            'error': error_msg,
            'answer': '抱歉，处理您的问题时出现了错误。'
        }


def get_response(messages: List[Dict[str, str]]) -> str:
    """
    获取LLM响应 - 保持原有接口
    """
    client = OpenAI(
        api_key="sk-",
        base_url="http://localhost:11434/v1",
    )

    completion = client.chat.completions.create(
        model="qwen:latest",
        messages=messages,
    )
    return completion.choices[0].message.content


def demo_optimized_rag():
    """演示优化后的RAG系统"""
    print("🎯 优化RAG系统演示")
    print("="*60)
    
    file_path = 'employee_manual.docx'
    
    try:
        # 1. 使用优化解析器加载文档
        print("1. 加载并解析员工手册...")
        if OPTIMIZED_PARSER_AVAILABLE:
            advanced_result = load_employee_file_advanced(file_path, include_metadata=True)
            if advanced_result['success']:
                print(f"✅ 优化解析成功:")
                print(f"   - 文本长度: {len(advanced_result['text'])} 字符")
                print(f"   - 处理时间: {advanced_result['processing_time']:.3f} 秒")
                if 'statistics' in advanced_result:
                    print(f"   - 元素统计: {advanced_result['statistics']}")
        
        # 2. 创建RAG友好的文档块
        print("\n2. 创建RAG文档块...")
        chunks = create_rag_chunks_from_file(file_path, chunk_size=400)
        print(f"✅ 生成 {len(chunks)} 个文档块")
        
        if chunks:
            avg_length = sum(len(chunk) for chunk in chunks) / len(chunks)
            print(f"   - 平均块长度: {avg_length:.0f} 字符")
        
        # 3. 存储到ChromaDB
        print("\n3. 存储到向量数据库...")
        if store_in_chroma_optimized(chunks):
            print("✅ 向量数据库存储成功")
        else:
            print("❌ 向量数据库存储失败")
            return
        
        # 4. 测试问答
        print("\n4. 测试RAG问答...")
        test_questions = [
            "员工工作满3年，年假有几天？",
            "公司的考勤制度是什么？",
            "试用期最长可以设置多久？"
        ]
        
        for question in test_questions:
            print(f"\n问题: {question}")
            result = rag_answer_question_enhanced(question)
            
            if result['success']:
                print(f"✅ 回答生成成功")
                print(f"   检索文档: {result['chunk_count']} 个")
                print(f"   回答: {result['answer'][:200]}...")
            else:
                print(f"❌ 回答生成失败: {result['error']}")
        
        print("\n🎉 优化RAG系统演示完成!")
        
    except Exception as e:
        print(f"❌ 演示过程出错: {str(e)}")


if __name__ == '__main__':
    # 运行演示
    demo_optimized_rag()
    
    # 也可以运行原有的逻辑进行对比
    print("\n" + "="*60)
    print("💡 使用说明:")
    print("1. 这个优化版本完全兼容原有的 rag1.py 接口")
    print("2. 新增了更多高级功能和更好的错误处理")
    print("3. 可以直接替换原有文件，或作为新的模块使用")
    print("4. 如果优化解析器不可用，会自动回退到原始实现")
    print("="*60)
